package com.bruce.modelscope.service.impl;

import java.util.concurrent.CountDownLatch;
import org.springframework.stereotype.Service;

import com.bruce.modelscope.service.ModelScopeService;

import dev.langchain4j.model.chat.ChatLanguageModel;
import dev.langchain4j.model.chat.StreamingChatLanguageModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import reactor.core.publisher.Flux;

@Service
@Slf4j
public class ModelScopeServiceImpl implements ModelScopeService {

    @Resource
    private ChatLanguageModel modelScopeChatModel;

    @Resource
    private StreamingChatLanguageModel modelScopeStreamingChatModel;

    @Override
    public String chat(String message) {
        log.info("message: {}", message);
        String result = modelScopeChatModel.chat(message);
        log.info("result: {}", result);
        return result;
    }

    @Override
    public Flux<String> chatStream(String message) {
        log.info("chatStream message: {}", message);
        return Flux.create(sink -> {
            // 用于等待响应完成的计数器
            CountDownLatch latch = new CountDownLatch(1);

            modelScopeStreamingChatModel.chat(message, new StreamingChatResponseHandler() {

                @Override
                public void onPartialResponse(String partialResponse) {
                    // 每次接收到的响应部分时，将部分响应部分发送给客户端
                    log.info("partialResponse: {}", partialResponse);
                    sink.next(partialResponse);
                }

                @Override
                public void onCompleteResponse(ChatResponse completeResponse) {
                    // 响应完成时，将完整响应发送给客户端
                    log.info("onCompleteResponse: " + completeResponse);
                    sink.complete();
                    latch.countDown();
                }

                @Override
                public void onError(Throwable error) {
                    // 错误处理
                    log.error("onError: {}", error);
                    sink.error(error);
                    latch.countDown();
                }
            });

            //阻塞直到响应完成或者发生错误
            try {
                latch.await();
            } catch (InterruptedException e) {
                log.error("InterruptedException: {}", e);
                Thread.currentThread().interrupt();
                sink.error(e);
            }
        });
    }

}
